# utils/network_monitor.py
import asyncio
import aiohttp
import time
from typing import Dict, List, Optional, Callable, Any
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class NetworkStatus(Enum):
    """网络状态枚举"""
    ONLINE = "online"
    OFFLINE = "offline"
    UNSTABLE = "unstable"
    SLOW = "slow"

@dataclass
class NetworkMetrics:
    """网络指标"""
    status: NetworkStatus
    latency: float  # 延迟（毫秒）
    last_check: datetime
    error_rate: float  # 错误率
    bandwidth: Optional[float] = None  # 带宽（Mbps）
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'status': self.status.value,
            'latency': self.latency,
            'last_check': self.last_check.isoformat(),
            'error_rate': self.error_rate,
            'bandwidth': self.bandwidth
        }

class NetworkMonitor:
    """网络监控器 - 监控网络状态并提供自动切换功能"""
    
    def __init__(
        self,
        check_interval: int = 30,  # 检查间隔（秒）
        timeout: float = 10.0,     # 请求超时时间
        stability_threshold: int = 5,  # 稳定性阈值（连续检查次数）
        latency_threshold: float = 2000.0  # 延迟阈值（毫秒）
    ):
        self.check_interval = check_interval
        self.timeout = timeout
        self.stability_threshold = stability_threshold
        self.latency_threshold = latency_threshold
        
        self.current_status = NetworkStatus.ONLINE
        self.metrics_history: List[NetworkMetrics] = []
        self.status_change_callbacks: List[Callable[[NetworkStatus], None]] = []
        
        self.monitoring_task: Optional[asyncio.Task] = None
        self.is_monitoring = False
        
        # 测试端点列表
        self.test_endpoints = [
            "https://www.google.com",
            "https://www.cloudflare.com",
            "https://httpbin.org/get",
            "https://api.github.com"
        ]
        
        # 状态统计
        self.consecutive_online = 0
        self.consecutive_offline = 0
        self.total_checks = 0
        self.error_count = 0
    
    async def start_monitoring(self):
        """开始网络监控"""
        if self.is_monitoring:
            logger.warning("网络监控已在运行中")
            return
        
        self.is_monitoring = True
        self.monitoring_task = asyncio.create_task(self._monitor_loop())
        logger.info("网络监控已启动")
    
    async def stop_monitoring(self):
        """停止网络监控"""
        if self.monitoring_task:
            self.monitoring_task.cancel()
            try:
                await self.monitoring_task
            except asyncio.CancelledError:
                pass
        
        self.is_monitoring = False
        logger.info("网络监控已停止")
    
    async def _monitor_loop(self):
        """监控循环"""
        while self.is_monitoring:
            try:
                # 执行网络检查
                metrics = await self._check_network()
                self._update_status(metrics)
                
                # 记录指标历史
                self.metrics_history.append(metrics)
                if len(self.metrics_history) > 100:  # 保留最近100次检查
                    self.metrics_history.pop(0)
                
                logger.debug(f"网络状态: {metrics.status.value}, 延迟: {metrics.latency:.2f}ms")
                
                # 等待下次检查
                await asyncio.sleep(self.check_interval)
                
            except Exception as e:
                logger.error(f"监控循环中发生错误: {e}")
                await asyncio.sleep(5)  # 错误后等待5秒再继续
    
    async def _check_network(self) -> NetworkMetrics:
        """检查网络状态"""
        start_time = time.time()
        self.total_checks += 1
        
        try:
            # 尝试多个测试端点
            latencies = []
            successful_checks = 0
            
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout)) as session:
                for endpoint in self.test_endpoints:
                    try:
                        check_start = time.time()
                        async with session.get(endpoint) as response:
                            latency = (time.time() - check_start) * 1000  # 转换为毫秒
                            latencies.append(latency)
                            successful_checks += 1
                            if response.status == 200:
                                break  # 有一个成功即可
                    except Exception:
                        continue  # 继续尝试下一个端点
            
            if successful_checks > 0:
                avg_latency = sum(latencies) / len(latencies)
                error_rate = (len(self.test_endpoints) - successful_checks) / len(self.test_endpoints)
            else:
                avg_latency = float('inf')
                error_rate = 1.0
                self.error_count += 1
        
        except Exception as e:
            logger.warning(f"网络检查失败: {e}")
            avg_latency = float('inf')
            error_rate = 1.0
            self.error_count += 1
        
        end_time = time.time()
        total_time = (end_time - start_time) * 1000  # 转换为毫秒
        
        # 确定网络状态
        if avg_latency == float('inf'):
            status = NetworkStatus.OFFLINE
        elif avg_latency > self.latency_threshold:
            status = NetworkStatus.SLOW
        elif error_rate > 0.5:  # 50%以上请求失败
            status = NetworkStatus.UNSTABLE
        else:
            status = NetworkStatus.ONLINE
        
        return NetworkMetrics(
            status=status,
            latency=avg_latency if avg_latency != float('inf') else 9999.0,
            last_check=datetime.now(),
            error_rate=error_rate
        )
    
    def _update_status(self, metrics: NetworkMetrics):
        """更新网络状态"""
        old_status = self.current_status
        self.current_status = metrics.status
        
        # 更新连续状态计数
        if self.current_status == NetworkStatus.ONLINE:
            self.consecutive_online += 1
            self.consecutive_offline = 0
        else:
            self.consecutive_offline += 1
            self.consecutive_online = 0
        
        # 如果状态发生变化，触发回调
        if old_status != self.current_status:
            logger.info(f"网络状态变化: {old_status.value} -> {self.current_status.value}")
            for callback in self.status_change_callbacks:
                try:
                    callback(self.current_status)
                except Exception as e:
                    logger.error(f"状态变化回调执行失败: {e}")
    
    def register_status_callback(self, callback: Callable[[NetworkStatus], None]):
        """注册状态变化回调"""
        self.status_change_callbacks.append(callback)
    
    def unregister_status_callback(self, callback: Callable[[NetworkStatus], None]):
        """注销状态变化回调"""
        if callback in self.status_change_callbacks:
            self.status_change_callbacks.remove(callback)
    
    def get_current_metrics(self) -> NetworkMetrics:
        """获取当前网络指标"""
        if self.metrics_history:
            return self.metrics_history[-1]
        else:
            return NetworkMetrics(
                status=self.current_status,
                latency=0.0,
                last_check=datetime.now(),
                error_rate=0.0
            )
    
    def get_status_history(self, limit: int = 20) -> List[NetworkMetrics]:
        """获取状态历史"""
        return self.metrics_history[-limit:] if self.metrics_history else []
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取统计信息"""
        if self.total_checks == 0:
            error_rate = 0.0
        else:
            error_rate = self.error_count / self.total_checks
        
        avg_latency = 0.0
        if self.metrics_history:
            valid_latencies = [m.latency for m in self.metrics_history if m.latency != 9999.0]
            if valid_latencies:
                avg_latency = sum(valid_latencies) / len(valid_latencies)
        
        return {
            'total_checks': self.total_checks,
            'error_count': self.error_count,
            'error_rate': error_rate,
            'avg_latency': avg_latency,
            'consecutive_online': self.consecutive_online,
            'consecutive_offline': self.consecutive_offline,
            'uptime_percentage': ((self.total_checks - self.error_count) / self.total_checks * 10) if self.total_checks > 0 else 100.0,
            'current_status': self.current_status.value
        }
    
    def should_use_online_model(self) -> bool:
        """判断是否应该使用在线模型"""
        if self.current_status == NetworkStatus.ONLINE:
            return True
        elif self.current_status == NetworkStatus.SLOW:
            # 对于慢速网络，根据延迟阈值决定
            current_metrics = self.get_current_metrics()
            return current_metrics.latency <= self.latency_threshold * 0.8  # 使用较低的阈值
        else:
            # 离线或不稳定时不使用在线模型
            return False
    
    async def test_model_access(self, model_service, test_prompt: str = "Hello") -> Dict[str, Any]:
        """测试模型访问能力"""
        if not hasattr(model_service, 'call_model'):
            return {
                'success': False,
                'error': 'Model service does not have call_model method',
                'latency': float('inf')
            }
        
        start_time = time.time()
        try:
            result = await model_service.call_model(
                prompt=test_prompt,
                max_tokens=10,
                temperature=0.1
            )
            
            latency = (time.time() - start_time) * 1000  # 转换为毫秒
            
            if result and result.get('success'):
                return {
                    'success': True,
                    'latency': latency,
                    'model_used': result.get('model', 'unknown')
                }
            else:
                return {
                    'success': False,
                    'error': 'Model call failed',
                    'latency': latency
                }
                
        except Exception as e:
            latency = (time.time() - start_time) * 1000
            return {
                'success': False,
                'error': str(e),
                'latency': latency
            }

class ModelFailoverManager:
    """模型故障转移管理器"""
    
    def __init__(self, network_monitor: NetworkMonitor):
        self.network_monitor = network_monitor
        self.online_model_service = None
        self.local_model_service = None
        self.fallback_enabled = True
        self.last_fallback_time = None
        
        # 注册网络状态变化回调
        self.network_monitor.register_status_callback(self._on_network_status_change)
    
    def set_online_model_service(self, service):
        """设置在线模型服务"""
        self.online_model_service = service
    
    def set_local_model_service(self, service):
        """设置本地模型服务"""
        self.local_model_service = service
    
    def _on_network_status_change(self, new_status: NetworkStatus):
        """网络状态变化回调"""
        logger.info(f"网络状态变化，当前状态: {new_status.value}")
    
    async def call_model_with_fallback(
        self,
        prompt: str,
        use_online_first: bool = True,
        **kwargs
    ) -> Dict[str, Any]:
        """带故障转移的模型调用"""
        if not self.fallback_enabled:
            # 不启用故障转移，直接使用指定服务
            if use_online_first and self.online_model_service:
                return await self.online_model_service.call_model(prompt, **kwargs)
            elif self.local_model_service:
                return await self.local_model_service.call_model(prompt, **kwargs)
            else:
                return {
                    'success': False,
                    'error': 'No model service available'
                }
        
        # 根据网络状态决定使用哪个模型
        if use_online_first and self.network_monitor.should_use_online_model() and self.online_model_service:
            # 首先尝试在线模型
            try:
                result = await self.online_model_service.call_model(prompt, **kwargs)
                if result and result.get('success'):
                    return result
                else:
                    logger.warning("在线模型调用失败，切换到本地模型")
            except Exception as e:
                logger.warning(f"在线模型调用异常: {e}，切换到本地模型")
        
        # 使用本地模型
        if self.local_model_service:
            try:
                result = await self.local_model_service.call_model(prompt, **kwargs)
                if result and result.get('success'):
                    # 记录故障转移时间
                    self.last_fallback_time = datetime.now()
                    return result
                else:
                    logger.error("本地模型调用也失败了")
                    return result
            except Exception as e:
                logger.error(f"本地模型调用异常: {e}")
                return {
                    'success': False,
                    'error': str(e)
                }
        else:
            return {
                'success': False,
                'error': 'Local model service not available'
            }
    
    def get_failover_status(self) -> Dict[str, Any]:
        """获取故障转移状态"""
        return {
            'fallback_enabled': self.fallback_enabled,
            'last_fallback_time': self.last_fallback_time.isoformat() if self.last_fallback_time else None,
            'should_use_online': self.network_monitor.should_use_online_model(),
            'current_network_status': self.network_monitor.current_status.value
        }

# 便利函数：创建默认网络监控器
def create_default_network_monitor(
    check_interval: int = 30,
    timeout: float = 10.0
) -> NetworkMonitor:
    """创建默认的网络监控器"""
    monitor = NetworkMonitor(
        check_interval=check_interval,
        timeout=timeout
    )
    return monitor

# 便利函数：创建默认故障转移管理器
def create_default_failover_manager(
    check_interval: int = 30,
    timeout: float = 10.0
) -> tuple[NetworkMonitor, ModelFailoverManager]:
    """创建默认的故障转移管理器"""
    network_monitor = create_default_network_monitor(check_interval, timeout)
    failover_manager = ModelFailoverManager(network_monitor)
    return network_monitor, failover_manager

# 使用示例
async def example_usage():
    """使用示例"""
    # 创建网络监控器
    monitor = create_default_network_monitor()
    
    # 开始监控
    await monitor.start_monitoring()
    
    # 获取当前指标
    metrics = monitor.get_current_metrics()
    print(f"当前网络状态: {metrics.status.value}")
    print(f"延迟: {metrics.latency:.2f}ms")
    
    # 获取统计信息
    stats = monitor.get_statistics()
    print(f"统计信息: {stats}")
    
    # 等待一段时间再查看
    await asyncio.sleep(35)  # 等待一次检查周期
    
    # 停止监控
    await monitor.stop_monitoring()

if __name__ == "__main__":
    # 注意：这个示例需要在异步环境中运行
    # asyncio.run(example_usage())
    pass
